import sys
#!/usr/bin/env python
"""导入文档到向量数据库的脚本"""

import os
import sys
import time
import argparse
from typing import List, Optional

# 设置项目根目录
def setup_project_root():
    """设置项目根目录为工作目录"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)  # 上一级目录是项目根目录
    os.chdir(project_root)
    sys.path.insert(0, project_root)
    return project_root

# 导入必要的模块
def import_required_modules():
    """导入必要的模块"""
    try:
        from backend.ingest import doc_ingester as document_ingester
        from backend.db import db_manager
        from backend.logging import get_logger
        return document_ingester, db_manager, get_logger
    except ImportError as e:
        print(f"导入模块失败: {str(e)}")
        print("请确保已安装所有依赖并正确设置了项目路径")
        sys.exit(1)

# 批量导入文档
def import_documents_from_directory(directory_path: str, recursive: bool = True, 
                                   supported_extensions: Optional[List[str]] = None):
    """从目录导入所有支持格式的文档"""
    # 导入必要的模块
    document_ingester, db_manager, get_logger = import_required_modules()
    logger = get_logger('document_import')
    
    try:
        # 初始化数据库连接
        logger.info("初始化数据库连接...")
        db_manager.init_db()
        
        # 文档摄取器已在导入时初始化完成
        logger.info("文档摄取器初始化完成")
        
        # 设置支持的文件格式
        if supported_extensions is None:
            supported_extensions = ['.txt', '.pdf', '.docx', '.doc', '.md', '.markdown', '.html', '.htm']
        
        # 确保目录存在
        if not os.path.exists(directory_path):
            raise FileNotFoundError(f"目录不存在: {directory_path}")
        
        # 统计信息
        total_files = 0
        successful_imports = 0
        failed_imports = 0
        start_time = time.time()
        
        logger.info(f"开始从目录导入文档: {directory_path}")
        logger.info(f"支持的文件格式: {', '.join(supported_extensions)}")
        logger.info(f"递归导入: {'是' if recursive else '否'}")
        
        # 遍历目录
        if recursive:
            for root, _, files in os.walk(directory_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    _, file_ext = os.path.splitext(file)
                    
                    # 只处理支持的文件格式
                    if file_ext.lower() in supported_extensions:
                        total_files += 1
                        try:
                            logger.info(f"正在处理文件 {total_files}: {file_path}")
                            success = document_ingester.add_document_to_vector_db(file_path)
                            if success:
                                successful_imports += 1
                                logger.info(f"成功导入文档: {file_path}")
                            else:
                                failed_imports += 1
                                logger.error(f"导入文档失败: {file_path}")
                        except Exception as e:
                            failed_imports += 1
                            logger.error(f"导入文档失败: {file_path}, 错误: {str(e)}")
        else:
            # 非递归模式，只处理当前目录
            for file in os.listdir(directory_path):
                file_path = os.path.join(directory_path, file)
                if os.path.isfile(file_path):
                    _, file_ext = os.path.splitext(file)
                    
                    # 只处理支持的文件格式
                    if file_ext.lower() in supported_extensions:
                        total_files += 1
                        try:
                            logger.info(f"正在处理文件 {total_files}: {file_path}")
                            document_ingester.ingest_document(file_path)
                            successful_imports += 1
                            logger.info(f"成功导入文档: {file_path}")
                        except Exception as e:
                            failed_imports += 1
                            logger.error(f"导入文档失败: {file_path}, 错误: {str(e)}")
        
        # 计算导入时间
        end_time = time.time()
        elapsed_time = end_time - start_time
        
        # 打印统计信息
        logger.info("===== 文档导入统计信息 ====")
        logger.info(f"总文件数: {total_files}")
        logger.info(f"成功导入: {successful_imports}")
        logger.info(f"导入失败: {failed_imports}")
        logger.info(f"总耗时: {elapsed_time:.2f} 秒")
        logger.info(f"平均每个文件耗时: {elapsed_time/total_files:.2f} 秒/文件" if total_files > 0 else "无文件处理")
        logger.info("=======================")
        
        # 注意：数据库当前没有提供文档统计功能
        logger.info("文档导入完成")
        
    except Exception as e:
        logger.error(f"导入过程失败: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

# 导入单个文档
def import_single_document(file_path: str):
    """导入单个文档"""
    # 导入必要的模块
    document_ingester, db_manager, get_logger = import_required_modules()
    logger = get_logger('document_import')
    
    try:
        # 初始化数据库连接
        logger.info("初始化数据库连接...")
        db_manager.init_db()
        
        # 初始化文档摄取器
        logger.info("初始化文档摄取器...")
        document_ingester.initialize()
        
        # 确保文件存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        logger.info(f"开始导入单个文档: {file_path}")
        start_time = time.time()
        
        # 导入文档
        document_ingester.ingest_document(file_path)
        
        # 计算导入时间
        end_time = time.time()
        elapsed_time = end_time - start_time
        
        logger.info(f"成功导入文档: {file_path}")
        logger.info(f"导入耗时: {elapsed_time:.2f} 秒")
        
        # 获取文档统计信息
        stats = db_manager.get_document_stats()
        logger.info(f"数据库中文档总数: {stats.get('document_count', 0)}")
        logger.info(f"数据库中向量总数: {stats.get('vector_count', 0)}")
        
    except Exception as e:
        logger.error(f"导入文档失败: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

# 主函数
def main():
    """主函数"""
    # 设置项目根目录
    project_root = setup_project_root()
    print(f"项目根目录: {project_root}")
    
    # 设置命令行参数
    parser = argparse.ArgumentParser(description="导入文档到向量数据库")
    parser.add_argument("--dir", type=str, help="要导入文档的目录路径")
    parser.add_argument("--file", type=str, help="要导入的单个文件路径")
    parser.add_argument("--recursive", action="store_true", default=True, help="是否递归导入子目录")
    parser.add_argument("--no-recursive", dest="recursive", action="store_false", help="不递归导入子目录")
    parser.add_argument("--format", type=str, nargs="*", help="支持的文件格式列表，例如: .txt .pdf .md")
    
    args = parser.parse_args()
    
    # 确保指定了目录或文件
    if not args.dir and not args.file:
        # 默认使用data/docs目录
        default_dir = os.path.join(project_root, 'data', 'docs')
        print(f"未指定目录或文件，将使用默认目录: {default_dir}")
        print("如果要指定目录，请使用 --dir 参数")
        print("如果要导入单个文件，请使用 --file 参数")
        
        # 确保默认目录存在
        if not os.path.exists(default_dir):
            print(f"默认目录不存在: {default_dir}")
            print("请创建该目录并添加文档，或使用 --dir 或 --file 参数指定路径")
            sys.exit(1)
        
        # 使用默认目录
        args.dir = default_dir
    
    # 执行导入
    if args.dir:
        supported_formats = args.format if args.format else None
        import_documents_from_directory(args.dir, args.recursive, supported_formats)
    elif args.file:
        import_single_document(args.file)
    
    print("文档导入完成")

# 执行主函数
if __name__ == "__main__":
    main()